// Copyright 2009 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Jeff Stewart, Gunnar Ritter
//
// Implements a probe that connects to a IMAP-over-SSL server and tests the
// ability to login and retrieve mail.  By default retrieves only one
// message, but can be optionally configured to fetch all messages.

#ifndef POPDIAG_PROBES_IMAPPROBE_H__
#define POPDIAG_PROBES_IMAPPROBE_H__

#include <QtNetwork/QHostInfo>
#include <QtNetwork/QTcpSocket>
#include <QtCore/QIODevice>
#include <QtCore/QSet>
#include <QtCore/QTextStream>
#include <QtCore/QTimer>
#include <QtCore/QTime>
#include <QtNetwork/QtNetwork>

#include "probes/linereader.h"
#include "probes/multiprobe.h"
#include "probes/probe.h"

namespace popdiag {

class ImapResponseReader;

typedef enum {
  IMAP_TOKEN_OTHER,  // e.g. a tag
  IMAP_TOKEN_NUMBER,
  IMAP_TOKEN_PLUS,
  IMAP_TOKEN_STAR,
  IMAP_TOKEN_OK,
  IMAP_TOKEN_PREAUTH,
  IMAP_TOKEN_NO,
  IMAP_TOKEN_BAD,
  IMAP_TOKEN_BYE,
  IMAP_TOKEN_CAPABILITY,
  IMAP_TOKEN_LBRACKET,
  IMAP_TOKEN_RBRACKET,
  IMAP_TOKEN_LPAREN,
  IMAP_TOKEN_RPAREN,
  IMAP_TOKEN_LANGLE,
  IMAP_TOKEN_RANGLE,
  IMAP_TOKEN_ALERT,
  IMAP_TOKEN_BADCHARSET,
  IMAP_TOKEN_PARSE,
  IMAP_TOKEN_PERMANENTFLAGS,
  IMAP_TOKEN_READ_ONLY,
  IMAP_TOKEN_READ_WRITE,
  IMAP_TOKEN_TRYCREATE,
  IMAP_TOKEN_UIDNEXT,
  IMAP_TOKEN_UIDVALIDITY,
  IMAP_TOKEN_UNSEEN,
  IMAP_TOKEN_FLAGS,
  IMAP_TOKEN_LIST,
  IMAP_TOKEN_LSUB,
  IMAP_TOKEN_SEARCH,
  IMAP_TOKEN_STATUS,
  IMAP_TOKEN_EXISTS,
  IMAP_TOKEN_RECENT,
  IMAP_TOKEN_EXPUNGE,
  IMAP_TOKEN_FETCH,
  IMAP_TOKEN_IMAP4,
  IMAP_TOKEN_IMAP4REV1,
  IMAP_TOKEN_STARTTLS,
  IMAP_TOKEN_AUTH_PLAIN,
  IMAP_TOKEN_LOGINDISABLED,
  IMAP_TOKEN_RFC822_SIZE,
  IMAP_TOKEN_UID,
  IMAP_TOKEN_BODY,
  IMAP_TOKEN_FLAG_ANSWERED,
  IMAP_TOKEN_FLAG_FLAGGED,
  IMAP_TOKEN_FLAG_DELETED,
  IMAP_TOKEN_FLAG_SEEN,
  IMAP_TOKEN_FLAG_DRAFT,
  IMAP_TOKEN_FLAG_RECENT,
  IMAP_TOKEN_STRING,
  IMAP_TOKEN_MALFORMED,
  IMAP_TOKEN_EOL
} ImapToken;


typedef enum {
  IMAP_RESPONSE_TYPE_CONTINUE,
  IMAP_RESPONSE_TYPE_TAGGED,
  IMAP_RESPONSE_TYPE_UNTAGGED,
  IMAP_RESPONSE_TYPE_FATAL,
  IMAP_RESPONSE_TYPE_UNKNOWN,
  IMAP_RESPONSE_TYPE_MALFORMED
} ImapResponseType;

typedef enum {
  IMAP_STATUS_RESPONSE_OK,
  IMAP_STATUS_RESPONSE_NO,
  IMAP_STATUS_RESPONSE_BAD,
  IMAP_STATUS_RESPONSE_PREAUTH,
  IMAP_STATUS_RESPONSE_BYE,
  IMAP_STATUS_RESPONSE_CAPABILITY,
  IMAP_STATUS_RESPONSE_EXPUNGE,
  IMAP_STATUS_RESPONSE_FETCH,
  IMAP_STATUS_RESPONSE_FLAGS,
  IMAP_STATUS_RESPONSE_LIST,
  IMAP_STATUS_RESPONSE_LSUB,
  IMAP_STATUS_RESPONSE_SEARCH,
  IMAP_STATUS_RESPONSE_STATUS,
  IMAP_STATUS_RESPONSE_EXISTS,
  IMAP_STATUS_RESPONSE_RECENT,
  IMAP_STATUS_RESPONSE_UNKNOWN,
  IMAP_STATUS_RESPONSE_MALFORMED,
  IMAP_STATUS_RESPONSE_NOT_APPLICABLE,
  IMAP_SOCKET_DROPPED,
  IMAP_STATUS_RESPONSE_TIMEOUT
} ImapStatusResponse;

typedef enum {
  IMAP_RESPONSE_CODE_NONE,
  IMAP_RESPONSE_CODE_ALERT,
  IMAP_RESPONSE_CODE_BADCHARSET,
  IMAP_RESPONSE_CODE_CAPABILITY,
  IMAP_RESPONSE_CODE_PARSE,
  IMAP_RESPONSE_CODE_PERMANENTFLAGS,
  IMAP_RESPONSE_CODE_READ_ONLY,
  IMAP_RESPONSE_CODE_READ_WRITE,
  IMAP_RESPONSE_CODE_TRYCREATE,
  IMAP_RESPONSE_CODE_UIDNEXT,
  IMAP_RESPONSE_CODE_UIDVALIDITY,
  IMAP_RESPONSE_CODE_UNSEEN,
  IMAP_RESPONSE_CODE_UNKNOWN,
  IMAP_RESPONSE_CODE_MALFORMED
} ImapResponseCode;

typedef enum {
  IMAP_CAPABILITY_IMAP4         = 0x00000001,
  IMAP_CAPABILITY_IMAP4REV1     = 0x00000002,
  IMAP_CAPABILITY_STARTTLS      = 0x00000004,
  IMAP_CAPABILITY_AUTH_PLAIN    = 0x00000008,
  IMAP_CAPABILITY_LOGINDISABLED = 0x00000010
} ImapCapability;

typedef enum {
  IMAP_FLAG_ANSWERED = 0x00000001,
  IMAP_FLAG_FLAGGED  = 0x00000002,
  IMAP_FLAG_DELETED  = 0x00000004,
  IMAP_FLAG_SEEN     = 0x00000008,
  IMAP_FLAG_DRAFT    = 0x00000010,
  IMAP_FLAG_RECENT   = 0x00000020
} ImapFlag;

// Spawns pop probes for each IP address detected by MultiProbe.
class ImapProbeFactory : public MultiProbeFactory {
 public:
  // Creates a factory that will merge these arguments with those of
  // CreateProbe and create ImapProbes on demand.
  ImapProbeFactory(const QString &username,
                  const QString &password,
                  const QString &imap_server_name,
                  bool pull_all_mail);

  // Merges these arguments with those of the constructor and returns a
  // constructed ImapProbe.
  virtual Probe *CreateProbe(const QHostAddress &addr, bool do_complete_probe);

 private:
  const QString username_;
  const QString password_;
  const QString imap_server_name_;
  const bool pull_all_mail_;
};

// Probes one IP address on port 993 talking IMAPS.  Probes only as far
// as banner verification unless do_complete_probe is true.  If pull_all_mail
// is true, we'll retrieve all mail.  Otherwise we'll pull the first message.
class ImapProbe : public Probe {
  Q_OBJECT
 public:
  // Creates a probe to probe imap over SSL on port 993.
  // username: The username to present.
  // password: The password to present.
  // imap_server_name: The DNS name of the IMAP server.  Used to verify the SSL
  //                  certificate.
  // pull_all_mail: If true, retrieve all mail.  Otherwise just the first msg.
  // do_complete_probe: If false, stop after banner validation.  Otherwise
  //                    do a full probe.
  ImapProbe(const QString &username,
           const QString &password,
           const QString &imap_server_name,
           const QHostAddress &host_address,
           bool pull_all_mail,
           bool do_complete_probe);
  virtual ~ImapProbe();

  // Fetch progress made in this probe.
  void GetProgress(int *progress, int *completed);

  // Called by the base class to start the probe
  void StartProbe();

  // Get a display string describing the results of this probe.
  QString GetDisplayText();

  // Retrieve a numeric representation of this fault.
  // 0 is reserved for "success".  Most other values are probe-specific.
  int GetFaultId();

  // Get a string representation of this probe's class.
  // Can be used with GetFaultId() to construct support urls, though
  // it is not currently used for this.
  inline QString GetProbeId() { return "imap"; }

  // Values in the common range must match multiprobe's values.
  typedef enum {
    // begin common range
    OK = 0,
    UNKNOWN = 1,
    RESOLVE_FAILED = 2,
    // end common range

    CONNECTION_REFUSED,
    CONNECT_TIMED_OUT,
    SERVER_CLOSED_PREMATURELY,
    SOCKET_ERROR,
    IMAP_SERVER_CAFILE_NOT_FOUND,
    IMAP_SERVER_CERT_SELFSIGNED,
    IMAP_SERVER_CERT_EXPIRED,
    IMAP_SERVER_CERT_NOTYETVALID,
    IMAP_SERVER_CERT_CANNOT_VALIDATE,
    IMAP_SERVER_CERT_NAME_MISMATCH,
    IMAP_SERVER_SHEDDING_LOAD,
    IMAP_SERVER_RESPONSE_MALFORMED,
    IMAP_RESPONSE_TIMED_OUT,
    IMAP_LOGIN_REJECTED,
    IMAP_LOGOUT_REJECTED,
    IMAP_LOGIN_FAILED_TEMP_SYSTEM_FAILURE,
    IMAP_NOT_ENABLED_FOR_USER,
    IMAP_NOT_ENABLED_FOR_DOMAIN,
    IMAP_NOT_AVAILABLE_FOR_USER,
    IMAP_WEB_LOGIN_REQUIRED,
    IMAP_ACCOUNT_NOT_READY,
    IMAP_LOGIN_FAILED_ABUSE,
    IMAP_LOGIN_FAILED_GEO_LOCKOUT,
    IMAP_LOGIN_FAILED_BAD_CREDS,
    IMAP_FETCH_STATS_FAILED,
    IMAP_FETCH_CONTENTS_FAILED,
    IMAP_SELECT_FAILED,
  } ImapProbeFault;

 private slots:
  void OnSslError(const QList<QSslError> &errors);
  void OnSslTimeout();
  void OnSslEstablished();
  void OnError(QAbstractSocket::SocketError err);

  void OnConnected();
  void OnDisconnected();
  void OnConnectTimedOut();
  void OnConnectError(QAbstractSocket::SocketError err);
  void OnCommandResponseReceived(ImapResponseReader *reader);

 private:
  typedef enum {
    INIT,
    CONNECT,
    NEGOTIATE_SSL,
    BANNER,
    CAPABILITY,
    LOGIN,
    SELECT,
    FETCH_STATS,
    FETCH_CONTENTS,
    LOGOUT
  } ProbeState;

  const char *StepNameFromState();
  ImapResponseReader *GetReader();

  void BeginSecureConversation();

  void ProcessOkResponse(ImapResponseReader *reader);
  void ProcessErrResponse(ImapResponseReader *reader);
  void ProcessExpungeResponse(ImapResponseReader *reader);
  void ProcessFetchResponse(ImapResponseReader *reader);

  // Do enhanced failure logging for a FETCH command.
  void LogFetchContentsFailure(ImapResponseReader *reader);

  void SetState(ProbeState t);
  void Cleanup();

  void ProbeDone(ImapProbeFault f);

  // These are called in sequence, with the invocation
  // of the next call placed from the corresponding
  // FinishProbeX above.  The exception is Connect, which
  // uses a variety of callbacks to terminate the various
  // error cases.
  void StartProbeConnect();
  void StartProbeNegotiateSsl();

  void StartProbeBanner();
  void BannerOk(ImapResponseReader *reader);
  void BannerRejected(ImapResponseReader *reader);
  void NatCheck(ImapResponseReader *banner_reader);

  void StartProbeCapability();
  void CapabilityOk(ImapResponseReader *reader);

  void StartProbeLogin();
  void LoginOk(ImapResponseReader *reader);
  void LoginRejected(ImapResponseReader *reader);

  void StartProbeSelect();
  void SelectOk(ImapResponseReader *reader);

  void StartProbeFetchStats();
  void FetchStatsOk(ImapResponseReader *reader);

  void StartProbeFetchContents();
  void FetchContentsOk(ImapResponseReader *reader);

  void ScheduleNextFetchContents();

  void StartProbeLogout();
  void LogoutOk(ImapResponseReader *reader);

  void SendCommand(const QString &cmd);

  static const int kTimeoutMs;
  const QString username_;
  const QString password_;
  const QString imap_server_;

  // The base object for the socket to use.
  scoped_ptr<RenamableSslSocket> base_socket_;

  // Read and write from here.
  RenamableSslSocket *socket_;

  // How far in our state machine have we gotten?
  ProbeState state_;

  // The highest value state_ will ever get.
  const ProbeState max_state_;

  // An enumerated value indicating the fault we've found.
  ImapProbeFault fault_;

  // If true, retrieve all mail during the IMAP session.
  // otherwise fetch the first message.
  bool pull_all_mail_;

  // If true, stop the probe when the banner is read.
  bool probe_banner_only_;

  // If true, we requested socket teardown.  Ignore socket errors.
  bool we_closed_;

  // How many bytes have we fetched?
  qulonglong bytes_retrieved_;

  // How many bytes do we expect we will fetch?
  qulonglong bytes_expected_;

  int msgs_in_mailbox_;
  int msgs_unseen_;
  int msgs_retrieved_;

  // The UIDs of the messages in the mailbox,
  // indexed by the 0-based message number (i.e. not IMAP starts-at-1 index).
  QList<qint64> uids_;

  // The lengths of the messages in the mailbox,
  // indexed by the 0-based message number (i.e. not IMAP starts-at-1 index).
  QList<qint64> lengths_;

  // The TCP address we're going to connect to in this probe.
  const QHostAddress target_address_;

  // Times out our connect attempts
  QTimer *connect_timer_;

  // Times out SSL negotiation.
  QTimer *ssl_negotiate_timer_;

  // Times our retrieves.
  QTime *fetch_timer_;

  // The current response reader, or NULL if we're not reading right now.
  ImapResponseReader *reader_;

  // Current tag number and resulting tag string.
  int tag_;
  QString tagstring_;

  // Whether we are (pre-)-authenticated.
  bool authenticated_;

  // Capability bits.
  int capabilities_;
};

// A word read from the server. The main point of this class is to do
// case-insensitive comparisons.
class ImapWord : public QString {
 public:
  ImapWord() : QString(NULL) { }
  explicit ImapWord(const QString &line) : QString(line) { }
  bool IsNumeric() const;
  bool operator==(const QString &other) const;
  bool operator!=(const QString &other) const;
  ImapWord operator=(const char *str);
  ImapWord operator=(QString str);
};

// An ImapLexer is a QString with the ability to retrieve tokens
// iteratively. The content of the string is the text that remains
// to be tokenized.
class ImapLexer : public QString {
 public:
  ImapLexer() : QString(NULL) { }
  explicit ImapLexer(const QString &line) : QString(line) { }
  // Get the next token.
  ImapToken Get();
  // Unget the current token. May not be called repeatedly.
  void Unget();
  // The literal value of the current token.
  const ImapWord lvalue() const { return lvalue_; }
  // The numeric value of the current token, if applicable.
  const qulonglong nvalue() const { return nvalue_; }

 private:
  ImapWord lvalue_;
  qint64 nvalue_;
};

// An ImapResponseReader takes over reading a response from the IMAP server.
// It will generate its Finished() signal when done, on error, or on timeout.
class ImapResponseReader : public QObject {
  Q_OBJECT
 public:
  // Create an object that can issue a command and read a IMAP response.
  // parent: takes memory ownership of this object.
  // src: the socket to read from
  // timeout_in_millis: If no response (line-wise) in this amount of time,
  //                    then declare the response timed out.
  ImapResponseReader(QObject *parent,
                     RenamableSslSocket *src,
                     int timeout_in_millis);
  ~ImapResponseReader();

  // Send a command and start reading the response.
  // The empty string is a special case.  It will start reading a response
  // without sending a command.
  void SendCommand(const QString &cmd, const QString &tag);

  // What kind of response did we get?
  ImapStatusResponse status_response() const { return status_response_; }

  // What kind of response did we get?
  ImapResponseType response_type() const { return response_type_; }

  // What response code did we get?
  ImapResponseCode response_code() const { return response_code_; }

  // Details of the response we got.
  const QString text() const { return text_; }
  const QString unknown() const { return unknown_; }
  qint64 number() const { return number_; }
  qint64 stringsize() const { return stringsize_; }
  qint64 uid() const { return uid_; }
  qint64 size() const { return size_; }
  int capabilities() const { return capabilities_; }
  int flags() const { return flags_; }

  // Describe the response for the diagnostic log.  Excerpt
  // any long response.
  QString describe();

  // The first response line.
  const QString first_response_line() const { return first_response_line_; }
  // Return the list of lines within a string or a literal.
  const QList<QString> &long_response_lines() const { return response_lines_; }
  // Get the tag string that we received.
  const QString tag_received() const { return tag_received_; }
  // This reader is no longer needed.
  void Cleanup();

 public slots:
  // Try to pull off a response line.
  void TryRead();

 signals:
  // This reader is done.  Use the accessors to dispatch the result.
  void ResponseRead(ImapResponseReader *me);

 private slots:
  void DataAvailable(LineReadResult r);

 private:
  // Finished reading a line.
  void Done();

  void ParseResponse(const QString &line);
  void ParseRespText();
  void ParseCapability();
  void ParseMailboxOrMessageData();
  void ParseFetch();
  void ParseFlags();

  void ParseLiteral();
  void ParseRestOfLiteral();
  void ParseLineAfterLiteral();

  RenamableSslSocket *socket_;
  QString first_response_line_;
  QList<QString> response_lines_;
  ImapResponseType response_type_;
  ImapStatusResponse status_response_;
  ImapResponseCode response_code_;
  int capabilities_;
  int flags_;
  QString text_;
  QString unknown_;
  qint64 number_;
  qint64 uid_;
  qint64 size_;
  qint64 stringsize_;

  LineReader reader_;
  ImapLexer lexer_;
  qint64 remaining_octets_in_literal_;
  bool processing_literal_;
  QString tag_received_;
  bool active_;
};

}  // namespace popdiag

#endif  // POPDIAG_PROBES_IMAPPROBE_H__
